import numpy as np
import matplotlib.pyplot as plt
from sklearn.cluster import DBSCAN
import matplotlib.patches as patches
import pickle
import pandas as pd
import openpyxl
import os

# 设置参数
center = (0, 0)  # 圆形区域中心坐标
radius = 20  # 圆形区域半径
current_dir = os.getcwd()

# 从文件中读取之前保存的 Matern 簇数据
file_path_Origin = os.path.join(current_dir, "matern_cluster_square_data.pkl")
with open(file_path_Origin, "rb") as f:
    points, parent_points_x, parent_points_y, square_markers = pickle.load(f)

# 为 square_indices 对应的点画正方形，并绘制同心圆
square_indices = np.where(square_markers == 1)[0]
square_points = points[square_indices]


# 创建一个包含两个子图的窗口
fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(20, 8))


distances = np.sqrt(
    (square_points[:, 0] - center[0]) ** 2 + (square_points[:, 1] - center[1]) ** 2
)
sorted_indices = np.argsort(distances)
# radius_w = distances[sorted_indices[0]]
radius_w = 5

print(radius_w)
# 确定GZ的圈
circle1 = plt.Circle(center, radius_w, fill=False, color="black")
ax1.add_patch(circle1)
circle2 = plt.Circle(center, radius_w, fill=False, color="black")
ax2.add_patch(circle2)


# 找出圆内的点
in_circle = (points[:, 0] - center[0]) ** 2 + (
    points[:, 1] - center[1]
) ** 2 <= radius_w**2


# 在第一个子图中绘制模拟 Matern 簇过程的图形
ax1.scatter(points[:, 0], points[:, 1], label="All Points")
parent_points = np.array(
    [(parent_points_x[i], parent_points_y[i]) for i in range(len(parent_points_x))]
)
ax1.scatter(
    parent_points[:, 0],
    parent_points[:, 1],
    marker="^",
    s=50,
    c="green",
    label="Cluster Centers",
)
ax1.set_xlim(center[0] - radius, center[0] + radius)
ax1.set_ylim(center[1] - radius, center[1] + radius)
ax2.set_xlim(center[0] - radius, center[0] + radius)
ax2.set_ylim(center[1] - radius, center[1] + radius)
ax1.set_aspect("equal", adjustable="box")
ax1.set_title("Matern Cluster Process")
ax1.set_xlabel("X-axis")
ax1.set_ylabel("Y-axis")
ax1.legend()


# 剔除圆内的点，获取用于聚类的点
points_for_clustering = points[~in_circle]


# 使用 DBSCAN 进行聚类
dbscan = DBSCAN(eps=2.7, min_samples=3)  # 这里的参数可根据实际情况调整
clusters = dbscan.fit_predict(points_for_clustering)


# 提取不同簇的点和噪声点，注意要根据 points_for_clustering 的索引对应回原 points 的索引来处理
unique_clusters = np.unique(clusters)


cluster_sets_for_draw = {cluster: [] for cluster in unique_clusters}


# 存储用于绘图的点
for i, cluster_label in enumerate(clusters):
    if cluster_label != -1:
        # 根据在 points_for_clustering 中的索引，找到在原 points 中的对应点
        index_in_original = np.where(~in_circle)[0][i]
        cluster_sets_for_draw[cluster_label].append(points[index_in_original])
    else:
        # 噪声点同样要找到在原 points 中的对应索引
        noise_indices = np.where(clusters == -1)[0]
        noise_indices_in_original = np.where(~in_circle)[0][noise_indices]
        for idx in noise_indices_in_original:
            cluster_sets_for_draw[-1].append(points[idx])


# 标注 Willie
ax1.scatter(0, 0, marker="+", s=100, c="r", label="Willie")
ax2.scatter(0, 0, marker="+", s=100, c="r", label="Willie")


# 将圆内的点变红
ax1.scatter(points[in_circle, 0], points[in_circle, 1], c="red")
ax2.scatter(points[in_circle, 0], points[in_circle, 1], c="red")


# 为每个簇绘制圆（新添加的逻辑）
for cluster_label, cluster_points in cluster_sets_for_draw.items():
    cluster_points = np.array(cluster_points)
    if cluster_label != -1:
        ax2.scatter(
            cluster_points[:, 0],
            cluster_points[:, 1],
            label=f"Cluster {cluster_label}",
        )
    else:
        ax2.scatter(
            cluster_points[:, 0],
            cluster_points[:, 1],
            label="Noise Points",
            marker="x",
            c="gray",
        )


# 存储信息的列表
info_list = []


# 统计每个簇内带正方形的点的数量，并添加标注
for cluster_label, cluster_points in cluster_sets_for_draw.items():
    cluster_points = np.array(cluster_points)
    if cluster_label != -1:
        # 找出在 points_for_clustering 中属于该簇的索引
        indices_in_cluster = np.where(clusters == cluster_label)[0]
        # 找出在原始 points 中对应的索引
        indices_in_original = np.where(~in_circle)[0][indices_in_cluster]
        # 找出属于该簇的 square_markers 标记
        cluster_square_markers = square_markers[indices_in_original]
        # 统计带正方形的点的数量
        num_square_points = np.sum(cluster_square_markers == 1)
        # 找出该簇中带正方形的点
        square_points_DB = cluster_points[cluster_square_markers == 1]
        if len(square_points_DB) > 0:
            # 计算带正方形的点到原点的距离
            # distances = np.sqrt(
            #     square_points_DB[:, 0] ** 2 + square_points_DB[:, 1] ** 2
            # )
            # # 找出最远的点的索引
            # min_index = np.argmin(distances)
            # min_point = square_points_DB[min_index]
            # # 计算距离
            # min_distance = distances[min_index]
            # # 绘制五角星
            # ax2.scatter(min_point[0], min_point[1], marker="*", s=200, c="blue")
            # 为带正方形的点绘制圆
            center_x = np.mean(square_points_DB[:, 0])
            center_y = np.mean(square_points_DB[:, 1])
            center_tuple = (center_x, center_y)
            distances = np.sqrt(
                (square_points_DB[:, 0] - center_x) ** 2
                + (square_points_DB[:, 1] - center_y) ** 2
            )
            max_index = np.argmax(distances)
            max_point = square_points_DB[max_index]
            max_distance = distances[max_index]
            radius = np.max(distances)
            ax2.scatter(max_point[0], max_point[1], marker="*", s=200, c="blue")
            circle = plt.Circle(center_tuple, radius, fill=False, color="blue")
            ax2.add_patch(circle)
            # 存储信息
            info_list.append(
                (
                    cluster_label,
                    max_point[0],
                    max_point[1],
                    center_x,
                    center_y,
                    num_square_points,
                    radius,
                )
            )
            ax2.text(
                center_x,
                center_y,
                str(num_square_points),
                fontsize=20,
                color="black",
                weight="bold",
            )


for cluster_label, cluster_points in cluster_sets_for_draw.items():
    cluster_points = np.array(cluster_points)
    if cluster_label != -1:
        # 找出在 points_for_clustering 中属于该簇的索引
        indices_in_cluster = np.where(clusters == cluster_label)[0]
        # 找出在原始 points 中对应的索引
        indices_in_original = np.where(~in_circle)[0][indices_in_cluster]
        # 找出属于该簇的 square_markers 标记
        cluster_square_markers = square_markers[indices_in_original]
        # 找出该簇中不带正方形的点
        no_square_points_DB = cluster_points[cluster_square_markers == 0]

        # 为不带正方形的点绘制圆
        if len(no_square_points_DB) > 0:
            center_x = np.mean(no_square_points_DB[:, 0])
            center_y = np.mean(no_square_points_DB[:, 1])
            center_tuple = (center_x, center_y)
            distances = np.sqrt(
                (no_square_points_DB[:, 0] - center_x) ** 2
                + (no_square_points_DB[:, 1] - center_y) ** 2
            )
            radius = np.max(distances)
            circle = plt.Circle(center_tuple, radius, fill=False, color="orange")
            ax2.add_patch(circle)


for x, y in square_points:
    # 创建新的 Rectangle 实例
    rect1 = patches.Rectangle(
        (x - 0.5, y - 0.5),
        1,
        1,
        linewidth=1,
        edgecolor="blue",
        facecolor="none",
    )
    ax1.add_patch(rect1)
    rect2 = patches.Rectangle(
        (x - 0.5, y - 0.5),
        1,
        1,
        linewidth=1,
        edgecolor="blue",
        facecolor="none",
    )
    ax2.add_patch(rect2)


# 计算所有带正方形的点到坐标原点的距离，并从大到小排序
distances = np.sqrt(square_points[:, 0] ** 2 + square_points[:, 1] ** 2)
sorted_distances = np.sort(distances)[::-1]


# 在 ax1 中画同心圆，圆心为坐标原点，半径为 sorted_distances 中的数据
for dist in sorted_distances:
    circle = plt.Circle(center, dist, fill=False, color="blue", linestyle="dashed")
    ax1.add_patch(circle)


ax2.set_aspect("equal", adjustable="box")
ax2.set_title("Matern Cluster Process with DBSCAN Clustering")
ax2.set_xlabel("X-axis")
ax2.set_ylabel("Y-axis")
# ax2.legend(loc="upper right", fontsize="small")


# 将信息存储到文件中

file_path = os.path.join(current_dir, "cluster_info.txt")
with open(file_path, "w") as f:
    for item in info_list:
        (
            cluster_label,
            max_point_x,
            max_point_y,
            num_square_points,
            center_x,
            center_y,
            radius,
        ) = item
        f.write(
            f"Cluster: {cluster_label}, Tar Point X: {max_point_x}, Tar Point Y: {max_point_y}, Num of Covert Points: {num_square_points}, ClusterCenterX:{center_x}, ClusterCenterY:{center_y}, DistanceToCovertClusterCenter: {radius}\n"
        )

# 将 info_list 中的数据存储为 Excel 表格
# 首先将 info_list 中的元组转换为列表，方便存储到 DataFrame 中
data_list = []
for item in info_list:
    (
        cluster_label,
        max_point_x,
        max_point_y,
        center_x,
        center_y,
        num_square_points,
        radius,
    ) = item
    # 将 min_point 转换为列表存储
    data_list.append(
        [
            cluster_label,
            max_point_x,
            max_point_y,
            center_x,
            center_y,
            num_square_points,
            radius,
        ]
    )
# 创建 DataFrame
df = pd.DataFrame(
    data_list,
    columns=[
        "Cluster Label",
        "Tar Point X",
        "Tar Point Y",
        "ClusterCenterX",
        "ClusterCenterY",
        "Number of Covert Points",
        "DistanceToCovertClusterCenter",
    ],
)
# 将 DataFrame 存储为 Excel 表格
file_path_Excel = os.path.join(current_dir, "cluster_info.xlsx")
df.to_excel(file_path_Excel, index=False)


# 提取中心圈内以及离群点中带正方形的点的坐标
center_circle_square_points = points[in_circle][square_markers[in_circle] == 1]
noise_square_points = []
if -1 in cluster_sets_for_draw:
    noise_points = np.array(cluster_sets_for_draw[-1])
    noise_square_indices = np.where(
        square_markers[np.isin(points, noise_points).all(axis=1)] == 1
    )[0]
    noise_square_points = points[noise_square_indices]


# 合并中心圈内和离群点中带正方形的点
all_square_points = np.vstack((center_circle_square_points, noise_square_points))


# 将这些点存储到新的 Excel 表中
if len(all_square_points) > 0:
    df_square = pd.DataFrame(all_square_points, columns=["X", "Y"])
    new_excel_path = os.path.join(current_dir, "special_square_points.xlsx")
    df_square.to_excel(new_excel_path, index=False)
else:
    print("No square points found in center circle or noise points.")


###########################
# 计算 AoI 的逻辑


# plt.show()
